from metagpt.team import Team
from test.utils import stream_jsonl, run_model, RESULT_DIR
import asyncio
import math
import os
import threading
import datetime
import shutil
import re
import subprocess
import time

# 最多线程数量，防止使用 API 并发数过高
MAX_THREAD_NUM = 5
pass_at_1 = 0.0
resume = 0
total_task_num = 0
start_time = 0
lock = threading.Lock()
lock2 = threading.Lock()
lock3 = threading.Lock()

# result_dir is like "./results/humanevalx_pass@1_15:00:00/"   
def run_test(task, rust_task, n: int, k: int, result_dir: str, use_naive=False, use_rag=True):
    global pass_at_1
    global total_task_num
    global start_time
    pass_time = 0
    rust_task_id = rust_task["task_id"].split("/")[-1]
    base_dir = result_dir + str(rust_task_id) + "/"
    os.makedirs(base_dir)
    cpp_file_name = base_dir + "main" + task["task_id"].replace('/', '_') + ".cpp"
    with open(cpp_file_name, "w") as f:
        f.write(task["prompt"] + task["canonical_solution"])
        
    with open(base_dir + rust_task["task_id"].replace('/', '_') + ".log", "w") as log:
        for i in range(n):
            # lock3.acquire(blocking=True)
            (code_text, compiler_message) =  asyncio.run(run_model(
                c_code=task["prompt"] + task["canonical_solution"],
                use_naive=use_naive,
                use_rag=use_rag,
                investment=5.0,
                max_round=20))
            # lock3.release()
            pattern = r"error"
            match = re.search(pattern, compiler_message, re.DOTALL)
            file_name = base_dir + "main" + rust_task["task_id"].replace('/', '_') + str(i) + "time" + ".rs"
            with open(file_name, "w") as f:
                f.write(code_text)
            if not match:
                # 通过编译
                pass_time += 1
                log.write("#" + str(i) + " Success\n")
            else:
                log.write("#" + str(i) + " Compile failed\n")
        # lock.acquire(blocking=True)
        value = 1.0 - (math.comb(n - pass_time, k) / math.comb(n, k))
        pass_at_1 *= total_task_num
        pass_at_1 += value
        total_task_num += 1
        pass_at_1 /= total_task_num
        log.write(f"Run time: {round(time.time() - start_time, 3)} s\n")
        log.write(f"pass@1: {pass_at_1}")
        # lock.release()
        
    
def test_humanevalx(k: int = 1, use_naive=False, use_rag=True):
    global start_time
    global resume
    n = 5
    
    # ths = []
    
    now = (datetime.datetime.now() + datetime.timedelta(hours=8)).strftime('%H:%M:%S')
    result_dir = RESULT_DIR + "humanevalx_pass@" + str(k) + "_" + now + "/"
    if os.path.exists(result_dir):
        shutil.rmtree(result_dir)
    os.makedirs(result_dir)
    start_time = time.time()
    j = 0
    for (task, rust_task) in zip(stream_jsonl("./testcases/humanevalx/humanevalx_cpp.jsonl.gz"), stream_jsonl("./testcases/humanevalx/humanevalx_rust.jsonl.gz")):
        if j < resume:
            j += 1
            continue
        run_test(task, rust_task, n, k, result_dir, use_naive, use_rag)
